import os
import requests
from datetime import datetime
# 通过 pip install 'volcengine-python-sdk[ark]' 安装方舟SDK
# pip install httpx[socks] 安装socksio库以及httpx的 SOCKS 扩展
# from volcenginesdkarkruntime import Ark
from openai import OpenAI

import json
from PIL import Image
from io import BytesIO


bozo = os.environ.get("BOZO") + "/Downloads/"
os.chdir(bozo)
# 定义尺寸映射字典，键为数字选项，值为对应的尺寸字符串
size_mapping = {
    1: "1536x2730"   # 对应 9:16 比例
    2: "2730x1536",  # 对应 16:9 比例
    3: "2048x2048",  # 对应 1:1 比例
}
size_option = int(input("请输入图片尺寸数字： (1➡️ 9:16、2➡️ 16:9、 3➡️ 1:1)\r\n "))
# 根据用户选择获取对应的尺寸值
selected_size = size_mapping[size_option]

tsc = input("请输入中文提示词：\r\n")

# client = Ark(   # 方舟SDK调用方式
client = OpenAI(
    # 此为默认路径，您可根据业务所在地域进行配置
    base_url="https://ark.cn-beijing.volces.com/api/v3",
    # 从环境变量中获取您的 API Key。
    api_key=os.environ.get("ARK_API_KEY"),
)

# imagesResponse = client.images.generate(
#     model="doubao-seedream-3-0-t2i-250415",
#     prompt="图片风格为「电影写真」，画面主体：一位身穿黄色连衣裙的女性，双手捧着一顶草编宽檐帽，站立于花海间。场景时空：户外田园场景，阳光明媚的晴天，背景有远山与开阔蓝天 。艺术风格：清新自然风，类似日系写真的清新感与生活气息。色彩色调：以明亮暖黄（花朵、衣裙）为主，搭配蓝天的澄澈蓝与草地的自然绿，色调清新明亮。构图视角：平视角度，人物居中，前景用大朵黄色花卉营造浅景深，突出主体又丰富画面层次。附加细节：花海中点缀多种黄色小花、少量其他色彩野花；远处山峦朦胧，天空有淡白云丝；人物眼镜增添文艺感，连衣裙的纽扣、袖口褶皱等细节清晰 。"
# )
# doubaoimg = imagesResponse.data[0].url

# 生成基于时间戳的文件名
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
filename = f"豆包_{timestamp}.jpeg"


response = client.images.generate(
    # 指定您创建的方舟推理接入点 ID
    model="doubao-seedream-3-0-t2i-250415",
    prompt=tsc,
    size=selected_size,
    response_format="url"        
)

doubaoimg = response.data[0].url
# print(f"图片URL地址: {doubaoimg}")

image = Image.open(BytesIO(requests.get(doubaoimg).content))
image.save(filename)

os.system('open '+ bozo+filename)
# os.system('open -a "/Applications/Doubao.app" ' + doubaoimg)
